from app.core.celery_config import celery_app
from app.core.database import SessionLocal
from app.models.task import Task, TaskStatus
from app.engine_manager import EngineManager
from app.storage import StorageManager
import json
import time
import logging
from datetime import datetime

logger = logging.getLogger(__name__)


@celery_app.task(bind=True)
async def run_fuzz_task(self, task_id: int):
    """执行模糊测试任务的Celery任务"""
    db = SessionLocal()
    task_manager = EngineManager()
    storage_manager = StorageManager()
    
    try:
        # 获取任务信息
        task = db.query(Task).filter(Task.id == task_id).first()
        if not task:
            logger.error(f"Task {task_id} not found")
            return {"status": "failed", "message": "Task not found"}
        
        # 更新任务状态
        task.status = TaskStatus.RUNNING
        task.started_at = datetime.now()
        db.commit()
        
        # 解析配置
        config = json.loads(task.config)
        
        # 执行模糊测试
        logger.info(f"Starting fuzz task {task_id} with engine {task.engine_name}")
        
        # 这里是模拟执行，实际应用中应该调用相应的引擎执行器
        result = await task_manager.execute_fuzz_task(
            engine_name=task.engine_name,
            target_id=task.target_id,
            config=config,
            task_id=task_id,
            update_progress=self.update_state
        )
        
        # 更新任务状态为完成
        task.status = TaskStatus.COMPLETED
        task.progress = 100.0
        task.completed_at = datetime.now()
        db.commit()
        
        logger.info(f"Fuzz task {task_id} completed successfully")
        return {"status": "completed", "result": result}
        
    except Exception as e:
        # 发生错误时更新任务状态
        task = db.query(Task).filter(Task.id == task_id).first()
        if task:
            task.status = TaskStatus.FAILED
            task.progress = 0.0
            db.commit()
        
        logger.error(f"Fuzz task {task_id} failed: {str(e)}")
        raise self.retry(exc=e, countdown=60, max_retries=3)
    
    finally:
        db.close()


@celery_app.task(bind=True)
async def update_task_progress(self, task_id: int, progress: float, status: str = None):
    """更新任务进度的任务"""
    db = SessionLocal()
    try:
        task = db.query(Task).filter(Task.id == task_id).first()
        if task:
            task.progress = progress
            if status:
                task.status = status
            db.commit()
            return {"status": "success"}
        return {"status": "failed", "message": "Task not found"}
    finally:
        db.close()


@celery_app.task
def cleanup_old_tasks(days: int = 30):
    """清理旧任务的定时任务"""
    db = SessionLocal()
    try:
        from datetime import datetime, timedelta
        cutoff_date = datetime.now() - timedelta(days=days)
        
        # 删除旧任务
        old_tasks = db.query(Task).filter(Task.created_at < cutoff_date).all()
        count = len(old_tasks)
        
        for task in old_tasks:
            db.delete(task)
        
        db.commit()
        logger.info(f"Cleaned up {count} old tasks")
        return {"status": "success", "count": count}
    finally:
        db.close()